# -*- coding: utf-8 -*-
# @Time    : 2022/10/10 22:31
# @Author  : 修远
# @File    : sqlalchemy.py
# @Comment :
from contextlib import contextmanager
from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy


class SQLAlchemy(_SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e
